import uuid
import time
import threading


# 雪花ID生成器类
class SnowflakeIdGenerator:
    """
    Twitter雪花ID算法实现
    64位ID结构:
    - 1位符号位 (0)
    - 41位时间戳 (毫秒)
    - 10位工作机器ID (5位数据中心ID + 5位工作节点ID)
    - 12位序列号
    """

    # 起始时间戳 (2024-01-01 00:00:00)
    START_TIMESTAMP = 1704067200000

    # 各部分的位数
    SEQUENCE_BITS = 12
    WORKER_ID_BITS = 10
    TIMESTAMP_BITS = 41

    # 最大值
    MAX_SEQUENCE = -1 ^ (-1 << SEQUENCE_BITS)
    MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)

    # 位移量
    WORKER_ID_SHIFT = SEQUENCE_BITS
    TIMESTAMP_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS

    def __init__(self, datacenter_id=1, worker_id=1):
        """
        初始化雪花ID生成器

        Args:
            datacenter_id: 数据中心ID (0-31)
            worker_id: 工作节点ID (0-31)
        """
        if datacenter_id > 31 or datacenter_id < 0:
            raise ValueError("datacenter_id must be between 0 and 31")
        if worker_id > 31 or worker_id < 0:
            raise ValueError("worker_id must be between 0 and 31")

        self.datacenter_id = datacenter_id
        self.worker_id = worker_id
        self.worker_id_val = (datacenter_id << 5) | worker_id
        self.sequence = 0
        self.last_timestamp = -1
        self.lock = threading.Lock()

    def _current_timestamp(self):
        """获取当前时间戳（毫秒）"""
        return int(time.time() * 1000)

    def _wait_next_millis(self, last_timestamp):
        """等待下一毫秒"""
        timestamp = self._current_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._current_timestamp()
        return timestamp

    def generate_id(self):
        """生成雪花ID"""
        with self.lock:
            timestamp = self._current_timestamp()

            if timestamp < self.last_timestamp:
                raise Exception("Clock moved backwards. Refusing to generate ID")

            if timestamp == self.last_timestamp:
                # 同一毫秒内，序列号自增
                self.sequence = (self.sequence + 1) & self.MAX_SEQUENCE
                if self.sequence == 0:
                    # 序列号溢出，等待下一毫秒
                    timestamp = self._wait_next_millis(self.last_timestamp)
            else:
                # 不同毫秒内，序列号重置为0
                self.sequence = 0

            self.last_timestamp = timestamp

            # 生成ID
            id_val = ((timestamp - self.START_TIMESTAMP) << self.TIMESTAMP_SHIFT) | \
                     (self.worker_id_val << self.WORKER_ID_SHIFT) | \
                     self.sequence

            return str(id_val)


# 全局雪花ID生成器实例
snowflake_generator = SnowflakeIdGenerator()


def generate_id():
    """生成UUID"""
    return str(uuid.uuid4())


def generate_snowflake_id():
    """生成雪花ID"""
    return snowflake_generator.generate_id()


def generate_short_id():
    """生成短ID（基于雪花ID的缩短版本）"""
    snowflake_id = snowflake_generator.generate_id()
    # 将雪花ID转换为更短的字符串（使用base62编码）
    return base62_encode(int(snowflake_id))


def base62_encode(num):
    """将数字转换为base62字符串"""
    if num == 0:
        return "0"

    chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    result = []

    while num > 0:
        result.append(chars[num % 62])
        num //= 62

    return "".join(reversed(result))


# 测试函数
if __name__ == "__main__":
    print("UUID:", generate_id())
    print("雪花ID:", generate_snowflake_id())
    print("短ID:", generate_short_id())

    # 测试生成多个ID
    print("\n生成10个雪花ID:")
    for i in range(100):
        print(f"ID {i + 1}: {generate_snowflake_id()}")
